import json
import os
from datetime import datetime
from typing import Dict, Any
import time


class EfficientStockTickSaver:
    """
    高效股票Tick实时保存器

    特点：
    - 针对单个Tick实时保存优化
    - 减少文件打开/关闭次数（使用文件缓冲）
    - 内存缓存最近处理的文件句柄
    - 异步写入支持（可选）
    - 自动处理日期变更
    """

    def __init__(self, base_path: str = "data", buffer_size: int = 10, flush_interval: float = 1.0):
        """
        初始化保存器

        参数:
        - base_path: 基础存储路径
        - buffer_size: 内存缓冲区大小（按股票代码）
        - flush_interval: 自动刷新间隔（秒）
        """
        self.base_path = base_path
        self.buffer_size = buffer_size
        self.flush_interval = flush_interval
        self._buffers = {}  # 股票代码 -> 数据缓冲区
        self._last_flush = time.time()
        os.makedirs(base_path, exist_ok=True)

    def _get_file_path(self, stock_code: str) -> str:
        """获取当日数据文件路径"""
        date_str = datetime.now().strftime("%Y-%m-%d")
        stock_dir = os.path.join(self.base_path, stock_code)
        os.makedirs(stock_dir, exist_ok=True)
        return os.path.join(stock_dir, f"{date_str}.json")

    def _load_existing_data(self, file_path: str) -> list:
        """加载已有数据"""
        if not os.path.exists(file_path):
            return []

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except:
            return []

    def _flush_buffer(self, stock_code: str):
        """刷新指定股票的缓冲区"""
        if stock_code not in self._buffers:
            return

        buffer = self._buffers[stock_code]
        if not buffer['data']:
            return

        file_path = self._get_file_path(stock_code)
        existing_data = self._load_existing_data(file_path)
        existing_data.extend(buffer['data'])

        # 使用临时文件写入模式避免写入过程中崩溃导致数据损坏
        temp_path = f"{file_path}.tmp"
        with open(temp_path, 'w', encoding='utf-8') as f:
            json.dump(existing_data, f, ensure_ascii=False, indent=4)

        # 原子性替换文件
        if os.path.exists(temp_path):
            os.replace(temp_path, file_path)

        buffer['data'] = []
        buffer['count'] = 0

    def _auto_flush(self):
        """定时自动刷新所有缓冲区"""
        current_time = time.time()
        if current_time - self._last_flush >= self.flush_interval:
            for code in list(self._buffers.keys()):
                self._flush_buffer(code)
            self._last_flush = current_time

    def save_tick(self,stock_code:str, tick_data: Dict[str, Any]) -> bool:
        """
        高效保存单个Tick数据

        参数:
        - tick_data: 必须包含"code"字段

        返回:
        - 是否成功接收（实际写入可能稍后发生）
        """

        # 初始化缓冲区
        if stock_code not in self._buffers:
            self._buffers[stock_code] = {
                'data': [],
                'count': 0,
                'last_date': datetime.now().strftime("%Y-%m-%d")
            }

        # 检查日期变更
        current_date = datetime.now().strftime("%Y-%m-%d")
        if self._buffers[stock_code]['last_date'] != current_date:
            self._flush_buffer(stock_code)
            self._buffers[stock_code]['last_date'] = current_date

        # 添加到缓冲区
        self._buffers[stock_code]['data'].append(tick_data)
        self._buffers[stock_code]['count'] += 1

        # 检查是否需要刷新
        if self._buffers[stock_code]['count'] >= self.buffer_size:
            self._flush_buffer(stock_code)

        # 定时自动刷新
        self._auto_flush()

        return True

    def flush_all(self):
        """立即刷新所有缓冲区"""
        for code in list(self._buffers.keys()):
            self._flush_buffer(code)
        self._last_flush = time.time()

    def __del__(self):
        """析构时自动刷新"""
        self.flush_all()

stock_saver = EfficientStockTickSaver(buffer_size=5, flush_interval=1.0)

# 使用示例
if __name__ == "__main__":
    # 创建保存器（每5条或每秒自动刷新）
    saver = EfficientStockTickSaver(buffer_size=5, flush_interval=1.0)

    # 模拟实时Tick数据
    ticks = [
        {"code": "600000", "price": 12.34, "volume": 100, "time": "09:30:00"},
        {"code": "600000", "price": 12.35, "volume": 200, "time": "09:30:01"},
        {"code": "600000", "price": 12.36, "volume": 150, "time": "09:30:02"},
        {"code": "600000", "price": 12.33, "volume": 300, "time": "09:30:03"},
        {"code": "600000", "price": 12.32, "volume": 250, "time": "09:30:04"},
        {"code": "600000", "price": 12.31, "volume": 180, "time": "09:30:05"},
        {"code": "000001", "price": 15.20, "volume": 100, "time": "09:30:00"},
        {"code": "000001", "price": 15.21, "volume": 200, "time": "09:30:01"},
    ]

    # 逐个保存
    for tick in ticks:
        saver.save_tick(tick)
        time.sleep(0.1)  # 模拟实时间隔

    # 确保所有数据写入文件
    saver.flush_all()